package com.hxl.tech.gateway.auth.oauth;


import com.hxl.tech.gateway.common.constant.AppConstant;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import java.util.HashMap;
import java.util.Map;


/**
 * 自定义JWT Token
 * @author soliddq
 * @date 2023-10-14
 */
@Component
public class MyJwtAccessTokenConverter extends JwtAccessTokenConverter {

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        DefaultOAuth2AccessToken result = new DefaultOAuth2AccessToken(accessToken);
        Map<String, Object> map = new HashMap<>(accessToken.getAdditionalInformation());
        Object principal = authentication.getPrincipal();

        RequestAttributes requestAttributes = RequestContextHolder.currentRequestAttributes();
        authentication.getOAuth2Request().getRequestParameters().entrySet().stream()
                .filter(s -> s.getKey().equals(AppConstant.HEADER_OPEN_MERCHANT_TYPE) || s.getKey().equals(AppConstant.HEADER_OPEN_MERCHANT_CODE))
                .forEach(s -> map.put(s.getKey(), s.getValue()));
        if(requestAttributes instanceof ServletRequestAttributes) {
            ((ServletRequestAttributes) requestAttributes).getRequest().getParameterMap().entrySet().stream().filter(e -> e.getKey().startsWith(AppConstant.HEADER_OPEN))
                    .forEach(e -> map.put(e.getKey(), e.getValue()[0]));
        }

        if(principal instanceof OAuthUser) {
            OAuthUser oAuthUser = (OAuthUser)authentication.getPrincipal();
            map.put(AppConstant.HEADER_OPEN_MERCHANT_TYPE, "Open-Api");
            map.put(AppConstant.HEADER_OPEN_MERCHANT_CODE, oAuthUser.getMerchantCode());
        }

        if(!authentication.getOAuth2Request().getExtensions().isEmpty()) {
            authentication.getOAuth2Request().getExtensions().entrySet().stream().filter(e -> e.getKey().startsWith(AppConstant.HEADER_OPEN))
                    .forEach(e -> map.put(e.getKey(), e.getValue()));
        }

        result.setAdditionalInformation(map);
        return super.enhance(result, authentication);
    }

    @Override
    public Map<String, Object> decode(String token) {
        return super.decode(token);
    }

    @Override
    public OAuth2Authentication extractAuthentication(Map<String, ?> map) {
        OAuth2Authentication authentication =  super.extractAuthentication(map);
        authentication.getOAuth2Request().getExtensions().putAll((Map)map);
        return authentication;
    }
}